/**
  ******************************************************************************
  * @file           : temp.cpp
  * @author         : wangyingjie
  * @brief          : None
  * @attention      : None
  * @date           : 2025/6/3
  ******************************************************************************
  */

#include <iostream>
#include <map>
#include <vector>
#include <deque>
#include <algorithm>
#include <iomanip>

#define HAVE_OPENCV_IMGPROC
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>

double INF = 10000;

int MAP_HEIGHT = 131, MAP_WIDTH = 131;

void applyRainbowColorMap(cv::Mat& src, cv::Mat& dst) {
    // 确保源图像是8位单通道
    CV_Assert(src.type() == CV_8UC1);

    // 创建正确的LUT：1x256像素，3通道，连续内存
    cv::Mat lut(1, 256, CV_8UC3);

    // 填充彩虹色映射
    for (int i = 0; i < 256; i++) {
        // 更精确的彩虹色映射
        cv::Vec3b color;

        if (i < 51) { // 紫色到蓝色
            color = cv::Vec3b(128 + i * 2, 0, 255);
        } else if (i < 102) { // 蓝色到青色
            color = cv::Vec3b(0, (i - 51) * 5, 255);
        } else if (i < 153) { // 青色到绿色
            color = cv::Vec3b(0, 255, 255 - (i - 102) * 5);
        } else if (i < 204) { // 绿色到黄色
            color = cv::Vec3b((i - 153) * 5, 255, 0);
        } else { // 黄色到红色
            color = cv::Vec3b(255, 255 - (i - 204) * 5, 0);
        }

        lut.at<cv::Vec3b>(0, i) = color;
    }

    // 应用LUT
    cv::LUT(src, lut, dst);
}

class Position2D
{
public:
    Position2D(int pos_x, int pos_y)
    {
        x = pos_x;
        y = pos_y;
    };

    int x = 0;
    int y = 0;
};

struct mapPoint
{
    double     occupancy    = INF;
    double     cost         = 0.0;
    double     exitCost     = 0.0;
    bool       isVisited    = false;
    bool       costComputed = false;
    bool       exitChecked  = false;
    bool       isTracked    = false;

    int        visitedNum   = 0;

    Position2D  next_point_position = Position2D(INF, INF);
    Position2D  prev_point_position = Position2D(INF, INF);
};


bool operator<(const Position2D &p1, const Position2D &p2) {
    return (p1.x < p2.x) || ((p1.x == p2.x) && (p1.y < p2.y));
}

bool operator == (const Position2D &p1, const Position2D &p2)
{
    return (p1.x == p2.x && p1.y == p2.y);
}


std::map<Position2D, mapPoint> prob_map;   // position-point

std::deque<Position2D> candidate_positions;

std::vector<Position2D> optimal_path;


std::vector<Position2D> getNeighbors(Position2D root_position)
{
    std::vector<Position2D> neighbor_positions = {
            Position2D(root_position.x-1  ,root_position.y-1),
            Position2D(root_position.x    ,root_position.y-1),
            Position2D(root_position.x+1  ,root_position.y-1),
            Position2D(root_position.x-1  ,root_position.y),

            Position2D(root_position.x+1  ,root_position.y),
            Position2D(root_position.x-1  ,root_position.y+1),
            Position2D(root_position.x    ,root_position.y+1),
            Position2D(root_position.x+1  ,root_position.y+1)
    };
    return neighbor_positions;
}


void updateNeighbors(Position2D root_position)
{
    std::vector<Position2D> neighbor_positions = getNeighbors(root_position);

    for(int i = 0; i < 8; i++)
    {
        if(!(0 <= neighbor_positions[i].x <= (MAP_HEIGHT-1) && 0 <= neighbor_positions[i].y <= (MAP_WIDTH-1)))
        {
            continue;
        }
        else if(prob_map[neighbor_positions[i]].occupancy != INF){

            if(!prob_map[neighbor_positions[i]].costComputed)
            {
                prob_map[neighbor_positions[i]].cost = prob_map[root_position].cost + 1;
                prob_map[neighbor_positions[i]].costComputed = true;
                candidate_positions.push_back(neighbor_positions[i]);
            }
        }
    }
}


void waveFrontSpread(Position2D start_position)
{
    prob_map[start_position].costComputed = true;
    Position2D curr_position = start_position;
    updateNeighbors(curr_position);

    while (!candidate_positions.empty())
    {
        curr_position = candidate_positions.front();
        candidate_positions.pop_front();
        updateNeighbors(curr_position);
    }
}

std::deque<Position2D> shortcut_candidate_positions;

std::deque<Position2D> sub_path;

bool isExitFound = false;
bool isPlanningFinished = false;


void resetExitCost()
{
    for(int i = 0; i < MAP_WIDTH; i++)
    {
        for (int j = 0; j < MAP_HEIGHT; j++)
        {
            prob_map[Position2D(i,j)].exitChecked = false;
            prob_map[Position2D(i,j)].exitCost = 0.0;
            prob_map[Position2D(i,j)].isTracked = false;
            shortcut_candidate_positions.clear();
            sub_path.clear();
        }
    }
}


Position2D findExit(Position2D stuck_position)
{
    std::vector<Position2D> neighbor_positions = getNeighbors(stuck_position);

    Position2D exit_point = Position2D(INF, INF);

    for(int i = 0; i < 8; i++)
    {
        if(!(0 <= neighbor_positions[i].x <= (MAP_HEIGHT-1) && 0 <= neighbor_positions[i].y <= (MAP_WIDTH-1)))
        {
            continue;
        }

        else if(prob_map[neighbor_positions[i]].occupancy != INF && !prob_map[neighbor_positions[i]].isVisited)
        {
            isExitFound  = true;
            prob_map[neighbor_positions[i]].prev_point_position = stuck_position;

            exit_point = neighbor_positions[i];
            break;
        }

        else if(prob_map[neighbor_positions[i]].occupancy != INF && !prob_map[neighbor_positions[i]].exitChecked)
        {
            prob_map[neighbor_positions[i]].exitChecked = true;
            prob_map[neighbor_positions[i]].exitCost = prob_map[stuck_position].exitCost + 1;
            prob_map[neighbor_positions[i]].prev_point_position = stuck_position;

            shortcut_candidate_positions.push_back(neighbor_positions[i]);
        }
    }

    return exit_point;
}



void cornerEscape(Position2D stuck_position)  // use dijkstra algorithm
{
    prob_map[stuck_position].exitChecked = true;
    prob_map[stuck_position].isTracked = true;

    Position2D curr_position = stuck_position;
    Position2D found_point = findExit(curr_position);   // won't find a exit in the first iteration since it is on the stuck point
//    std::cout << "shortcut_candidate_positions:" << shortcut_candidate_positions.size() << std::endl;

    while(!shortcut_candidate_positions.empty())
    {
        curr_position = shortcut_candidate_positions.front();
        shortcut_candidate_positions.pop_front();

        found_point = findExit(curr_position);

        if(found_point.x != INF && found_point.y != INF)
        {


            for(Position2D position = found_point;
                prob_map[position].prev_point_position.x != INF && prob_map[position].prev_point_position.y != INF, !prob_map[position].isTracked;
                position = prob_map[position].prev_point_position)
            {
                prob_map[position].isVisited = true;
                prob_map[position].visitedNum++;
                prob_map[position].isTracked = true;

                sub_path.push_front(position);
            }


            // for debugging
//            for(int i = 0; i < sub_path.size(); i++)
//            {
//                std::cout << sub_path[i].x << ", " << sub_path[i].y << std::endl;
//            }

            optimal_path.insert(optimal_path.end(), sub_path.begin(), sub_path.end());

            isExitFound = false; // reset the indicator
            resetExitCost();
            return;
        }

    }

    isPlanningFinished = true;
    std::cout << "path planning finished." << std::endl;
}


void findNextStep(Position2D position)
{
    bool isStuck = true;
    double min_cost = INF;

    Position2D next_position = position;

    std::vector<Position2D> neighbor_positions = getNeighbors(position);

    for(int i = 0; i < 8; i++)
    {
        if(!(0 <= neighbor_positions[i].x <= (MAP_HEIGHT-1) && 0 <= neighbor_positions[i].y <= (MAP_WIDTH-1)))
        {
            continue;
        }
        else if((prob_map[neighbor_positions[i]].occupancy != INF) && (!prob_map[neighbor_positions[i]].isVisited))
        {
            if(prob_map[neighbor_positions[i]].cost < min_cost)
            {
                min_cost = prob_map[neighbor_positions[i]].cost;
                next_position = neighbor_positions[i];

                isStuck = false;
            }
        }
    }
    if(isStuck)
    {
//        std::cout << "stuck point: " << position.x << ", " << position.y << std::endl;
        cornerEscape(position);
    }
    else
    {
        prob_map[next_position].isVisited = true;
        prob_map[next_position].visitedNum++;
        optimal_path.push_back(next_position);
//        std::cout << next_position.x << ", " << next_position.y << std::endl;
    }
}


void completeCoveragePathPlanning(Position2D start_position)
{
    waveFrontSpread(start_position);

    prob_map[start_position].isVisited = true;
    prob_map[start_position].visitedNum++;
    optimal_path.push_back(start_position);

    Position2D curr_position = start_position;

    while(!isPlanningFinished)
    {
        findNextStep(curr_position);
        curr_position = optimal_path.back();
    }
}


int main() {
    // 构建一个简单的地图
    for(int i = 0; i < MAP_WIDTH; i++)
    {
        for (int j = 0; j < MAP_HEIGHT; j++)
        {
            if(i != 0 && i != (MAP_WIDTH-1) && j != 0 && j != (MAP_HEIGHT-1))
            {
                prob_map[Position2D(i,j)].occupancy = 1.0;
            }
            else
            {
                prob_map[Position2D(i,j)].occupancy = INF;
                prob_map[Position2D(i,j)].cost = INF;
            }
        }
    }

    // 向地图中添加障碍物
    for(int i = 40; i < 90; i++)
    {
        for (int j = 40; j < 90; j++)
        {
            if(i == 40 || i == 89 || j == 89)
            {
                prob_map[Position2D(i,j)].occupancy = INF;
                prob_map[Position2D(i,j)].cost = INF;
            }
        }
    }

    // 定义路径规划的起点
    Position2D start = Position2D(41,88);

    // 从起点开始执行全覆盖路径规划
    completeCoveragePathPlanning(start);

    // 输出换行符
    std::cout << std::endl << std::endl;

    // 初始化成本地图
    cv::Mat cost_map = cv::Mat::zeros(MAP_HEIGHT, MAP_WIDTH, CV_32FC1);
    for(int i = 0; i < MAP_WIDTH; i++)
    {
        for (int j = 0; j < MAP_HEIGHT; j++)
        {
            if(i != 0 && i != (MAP_HEIGHT-1) && j != 0 && j != (MAP_WIDTH-1))
            {
                cost_map.at<float>(j,i) = prob_map[Position2D(i,j)].cost*1.2;
            }
            else
            {
                cost_map.at<float>(j,i) = 255.0;
            }
        }
    }
    for(int i = 40; i < 90; i++)
    {
        for (int j = 40; j < 90; j++)
        {
            if(i == 40 || i == 89 || j == 89)
            {
                cost_map.at<float>(j,i) = 255.0;
            }
        }
    }
    // 将成本地图转换为8位单通道图像
    cost_map.convertTo(cost_map, CV_8UC1);
    // 对成本地图应用颜色映射
     // cv::applyColorMap(cost_map, cost_map, cv::COLORMAP_RAINBOW);
    // 显示成本地图
    cv::namedWindow("cost map", cv::WINDOW_NORMAL);
     cv::imshow("cost map", cost_map);
    cv::waitKey(0);

    // 初始化用于显示路径的地图
    cv::Mat map = cv::Mat::zeros(MAP_HEIGHT, MAP_WIDTH, CV_32FC3);
    for(int i = 0; i < MAP_WIDTH; i++)
    {
        for (int j = 0; j < MAP_HEIGHT; j++)
        {
            if(i != 0 && i != (MAP_HEIGHT-1) && j != 0 && j != (MAP_WIDTH-1))
            {
                map.at<cv::Vec3f>(j,i) = cv::Vec3f(255.0, 255.0, 255.0);
            }
            else
            {
                map.at<cv::Vec3f>(j,i) = cv::Vec3f(0.0, 0.0, 0.0);
            }
        }
    }
    for(int i = 40; i < 90; i++)
    {
        for (int j = 40; j < 90; j++)
        {
            if(i == 40 || i == 89 || j == 89)
            {
                map.at<cv::Vec3f>(j,i) = cv::Vec3f(0.0, 0.0, 0.0);
            }
            else
            {
                map.at<cv::Vec3f>(j,i) = cv::Vec3f(255.0, 255.0, 255.0);
            }
        }
    }

    // 初始化热力图以显示路径经过频率
    cv::Mat hotmap = cv::Mat::zeros(map.rows, map.cols, CV_32FC1);

    // 遍历并显示最优路径
    std::vector<Position2D>::iterator it = optimal_path.begin();
    cv::namedWindow("trajectory", cv::WINDOW_NORMAL);
    while(it != optimal_path.end())
    {
        map.at<cv::Vec3f>(it->y, it->x) = cv::Vec3f(0.0, 0.0, 255.0);
        hotmap.at<float>(it->y, it->x) += 51.0;
        cv::imshow("trajectory", map);
        cv::waitKey(1);
        it++;
    }
    cv::waitKey(0);

    // 将热力图转换为8位单通道图像
    hotmap.convertTo(hotmap, CV_8UC1);
    // 对热力图应用颜色映射
     // cv::applyColorMap(hotmap, hotmap, cv::COLORMAP_HOT);
    // 显示热力图
    cv::namedWindow("colorHotmap", cv::WINDOW_NORMAL);
     cv::imshow("colorHotmap", hotmap);
    cv::waitKey(0);

    return 0;
}
